{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction\n",
    "\n",
    "### In this notebook we use [SRGAN](https://arxiv.org/abs/1609.04802) to perform Super Resolution on [CelebA dataset](http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://miro.medium.com/max/1000/1*-txs0CYANMq5CVi0tp5DCg.png\" width=\"900\" height=\"900\"/>\n",
    "<h4></h4>\n",
    "<h4><center>Image Source:  <a href=\"https://arxiv.org/abs/1609.04802\">Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network [C. Ledig et al.]</a></center></h4>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Libraries 📚⬇"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os, math, sys\n",
    "import glob, itertools\n",
    "import argparse, random\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.autograd import Variable\n",
    "from torchvision.models import vgg19\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "from torchvision.utils import save_image, make_grid\n",
    "\n",
    "import plotly\n",
    "import plotly.express as px\n",
    "import plotly.graph_objects as go\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from PIL import Image\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "random.seed(42)\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings ⚙️"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load pretrained models\n",
    "load_pretrained_models = True\n",
    "# number of epochs of training\n",
    "n_epochs = 2\n",
    "# name of the dataset\n",
    "dataset_path = \"celebA\"\n",
    "# size of the batches\n",
    "batch_size = 16\n",
    "# adam: learning rate\n",
    "lr = 0.00001\n",
    "# adam: decay of first order momentum of gradient\n",
    "b1 = 0.9\n",
    "# adam: decay of second order momentum of gradient\n",
    "b2 = 0.999\n",
    "# epoch from which to start lr decay\n",
    "decay_epoch = 100\n",
    "# number of cpu threads to use during batch generation\n",
    "n_cpu = 8\n",
    "# high res. image height\n",
    "hr_height = 256\n",
    "# high res. image width\n",
    "hr_width = 256\n",
    "# number of image channels\n",
    "channels = 3\n",
    "\n",
    "os.makedirs(\"image\", exist_ok=True)\n",
    "os.makedirs(\"saved_model\", exist_ok=True)\n",
    "\n",
    "\n",
    "hr_shape = (hr_height, hr_width)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Dataset Class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Normalization parameters for pre-trained PyTorch models\n",
    "\n",
    "mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)\n",
    "std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)\n",
    "\n",
    "class ImageDataset(Dataset):\n",
    "    def __init__(self, files, hr_shape):\n",
    "        hr_height, hr_width = hr_shape\n",
    "        # Transforms for low resolution images and high resolution images\n",
    "        self.lr_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.Resize((hr_height // 4, hr_height // 4), Image.BICUBIC),\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize(mean, std),\n",
    "            ]\n",
    "        )\n",
    "        \n",
    "        self.hr_transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.Resize((hr_height, hr_height), Image.BICUBIC),\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize(mean, std),\n",
    "            ]\n",
    "        )\n",
    "        self.files = files\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        img = Image.open(self.files[index % len(self.files)]).convert('RGB')\n",
    "        img_lr = self.lr_transform(img)\n",
    "        img_hr = self.hr_transform(img)\n",
    "\n",
    "\n",
    "        return {\"lr\": img_lr, \"hr\": img_hr}\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.files)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get Train/Test Dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_paths, test_paths = train_test_split(sorted(glob.glob(dataset_path + \"/*.*\")), test_size=0.02, random_state=42)\n",
    "train_paths=train_paths[0:800]\n",
    "train_paths=train_paths[0:200]\n",
    "train_dataloader = DataLoader(ImageDataset(train_paths, hr_shape=hr_shape), batch_size=batch_size, shuffle=True)\n",
    "test_dataloader = DataLoader(ImageDataset(test_paths, hr_shape=hr_shape), batch_size=int(batch_size*0.75), shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "600\n"
     ]
    }
   ],
   "source": [
    "print(len(test_paths))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<h3><center>Model Architecture</center></h3>\n",
    "<img src=\"https://miro.medium.com/max/1000/1*zsiBj3IL4ALeLgsCeQ3lyA.png\" width=\"900\" height=\"900\"/>\n",
    "<h4></h4>\n",
    "<h4><center>Image Source:  <a href=\"https://arxiv.org/abs/1609.04802\">Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network [C. Ledig et al.]</a></center></h4>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Model Classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VGGFeatureExtractor(nn.Module):\n",
    "    def __init__(self, feature_layer=34, use_bn=False, use_input_norm=True,\n",
    "                 device=torch.device('mps')):\n",
    "        super(VGGFeatureExtractor, self).__init__()\n",
    "        self.use_input_norm = use_input_norm\n",
    "        if use_bn:\n",
    "            model = torchvision.models.vgg19_bn(pretrained=True)\n",
    "        else:\n",
    "            model = torchvision.models.vgg19(pretrained=False)\n",
    "        if self.use_input_norm:\n",
    "            mean = torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)\n",
    "            # [0.485 - 1, 0.456 - 1, 0.406 - 1] if input in range [-1, 1]\n",
    "            std = torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)\n",
    "            # [0.229 * 2, 0.224 * 2, 0.225 * 2] if input in range [-1, 1]\n",
    "            self.register_buffer('mean', mean) # 构建处理步骤的名字\n",
    "            self.register_buffer('std', std)\n",
    "        self.features = nn.Sequential(*list(model.features.children())[:(feature_layer + 1)])\n",
    "        # No need to BP to variable\n",
    "        for k, v in self.features.named_parameters():\n",
    "            v.requires_grad = False # 为了防止梯度进行更新\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Assume input range is [0, 1]\n",
    "        if self.use_input_norm:\n",
    "            x = (x - self.mean) / self.std\n",
    "        output = self.features(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting models\n",
      "  Downloading models-0.9.3.tar.gz (16 kB)\n",
      "  Preparing metadata (setup.py) ... \u001b[?25lerror\n",
      "  \u001b[1;31merror\u001b[0m: \u001b[1msubprocess-exited-with-error\u001b[0m\n",
      "  \n",
      "  \u001b[31m×\u001b[0m \u001b[32mpython setup.py egg_info\u001b[0m did not run successfully.\n",
      "  \u001b[31m│\u001b[0m exit code: \u001b[1;36m1\u001b[0m\n",
      "  \u001b[31m╰─>\u001b[0m \u001b[31m[8 lines of output]\u001b[0m\n",
      "  \u001b[31m   \u001b[0m Traceback (most recent call last):\n",
      "  \u001b[31m   \u001b[0m   File \"<string>\", line 2, in <module>\n",
      "  \u001b[31m   \u001b[0m   File \"<pip-setuptools-caller>\", line 34, in <module>\n",
      "  \u001b[31m   \u001b[0m   File \"/private/var/folders/j3/xgtp4_0s2bvdwk75_tfv10mc0000gn/T/pip-install-zp9bqt39/models_38657eb1844a42b2a1f425048038d392/setup.py\", line 25, in <module>\n",
      "  \u001b[31m   \u001b[0m     import models\n",
      "  \u001b[31m   \u001b[0m   File \"/private/var/folders/j3/xgtp4_0s2bvdwk75_tfv10mc0000gn/T/pip-install-zp9bqt39/models_38657eb1844a42b2a1f425048038d392/models/__init__.py\", line 23, in <module>\n",
      "  \u001b[31m   \u001b[0m     from base import *\n",
      "  \u001b[31m   \u001b[0m ModuleNotFoundError: No module named 'base'\n",
      "  \u001b[31m   \u001b[0m \u001b[31m[end of output]\u001b[0m\n",
      "  \n",
      "  \u001b[1;35mnote\u001b[0m: This error originates from a subprocess, and is likely not a problem with pip.\n",
      "\u001b[1;31merror\u001b[0m: \u001b[1mmetadata-generation-failed\u001b[0m\n",
      "\n",
      "\u001b[31m×\u001b[0m Encountered error while generating package metadata.\n",
      "\u001b[31m╰─>\u001b[0m See above for output.\n",
      "\n",
      "\u001b[1;35mnote\u001b[0m: This is an issue with the package mentioned above, not pip.\n",
      "\u001b[1;36mhint\u001b[0m: See above for details.\n",
      "\u001b[?25h"
     ]
    },
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'models'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[21], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnn\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfunctional\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mF\u001b[39;00m\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mmodels\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodules\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodule_util\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mmutil\u001b[39;00m\n\u001b[1;32m      9\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mResidualDenseBlock_5C\u001b[39;00m(nn\u001b[38;5;241m.\u001b[39mModule):\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\u001b[38;5;28mself\u001b[39m, nf\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m, gc\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m32\u001b[39m, bias\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m):\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'models'"
     ]
    }
   ],
   "source": [
    "import functools\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "\n",
    "class ResidualDenseBlock_5C(nn.Module):\n",
    "    def __init__(self, nf=64, gc=32, bias=True):\n",
    "        super(ResidualDenseBlock_5C, self).__init__()\n",
    "        # gc: growth channel, i.e. intermediate channels\n",
    "        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)\n",
    "        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)\n",
    "        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n",
    "\n",
    "        # initialization\n",
    "        mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x1 = self.lrelu(self.conv1(x))\n",
    "        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))\n",
    "        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))\n",
    "        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))\n",
    "        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))\n",
    "        return x5 * 0.2 + x\n",
    "\n",
    "\n",
    "class RRDB(nn.Module):\n",
    "    '''Residual in Residual Dense Block'''\n",
    "\n",
    "    def __init__(self, nf, gc=32):\n",
    "        super(RRDB, self).__init__()\n",
    "        self.RDB1 = ResidualDenseBlock_5C(nf, gc)\n",
    "        self.RDB2 = ResidualDenseBlock_5C(nf, gc)\n",
    "        self.RDB3 = ResidualDenseBlock_5C(nf, gc)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.RDB1(x)\n",
    "        out = self.RDB2(out)\n",
    "        out = self.RDB3(out)\n",
    "        return out * 0.2 + x\n",
    "\n",
    "\n",
    "class RRDBNet(nn.Module):\n",
    "    def __init__(self, in_nc, out_nc, nf, nb, gc=32):\n",
    "        super(RRDBNet, self).__init__()\n",
    "        RRDB_block_f = functools.partial(RRDB, nf=nf, gc=gc)\n",
    "\n",
    "        self.conv_first = nn.Conv2d(in_nc, nf, 3, 1, 1, bias=True) # 进行第一次的卷积\n",
    "        self.RRDB_trunk = mutil.make_layer(RRDB_block_f, nb) # 添加RRDB残差块\n",
    "        self.trunk_conv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True) # 再次进行卷积操作\n",
    "        #### upsampling\n",
    "        self.upconv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.upconv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.HRconv = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n",
    "        self.conv_last = nn.Conv2d(nf, out_nc, 3, 1, 1, bias=True)\n",
    "\n",
    "        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        fea = self.conv_first(x)\n",
    "        trunk = self.trunk_conv(self.RRDB_trunk(fea))\n",
    "        fea = fea + trunk\n",
    "\n",
    "        fea = self.lrelu(self.upconv1(F.interpolate(fea, scale_factor=2, mode='nearest')))\n",
    "        fea = self.lrelu(self.upconv2(F.interpolate(fea, scale_factor=2, mode='nearest')))\n",
    "        out = self.conv_last(self.lrelu(self.HRconv(fea)))\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'RRDBNet' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[19], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m model\u001b[38;5;241m=\u001b[39m\u001b[43mRRDBNet\u001b[49m()\n",
      "\u001b[0;31mNameError\u001b[0m: name 'RRDBNet' is not defined"
     ]
    }
   ],
   "source": [
    "model=RRDBNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FeatureExtractor(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(FeatureExtractor, self).__init__()\n",
    "        vgg19_model = vgg19(pretrained=True)\n",
    "        self.feature_extractor = nn.Sequential(*list(vgg19_model.features.children())[:18])\n",
    "\n",
    "    def forward(self, img):\n",
    "        return self.feature_extractor(img)\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_features):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv_block = nn.Sequential(\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            #nn.BatchNorm2d(in_features, 0.8),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            #nn.BatchNorm2d(in_features, 0.8),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            #nn.BatchNorm2d(in_features, 0.8),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Conv2d(in_features, in_features, kernel_size=3, stride=1, padding=1),\n",
    "            #nn.BatchNorm2d(in_features, 0.8),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.conv_block(x)*0.8\n",
    "\n",
    "\n",
    "class GeneratorResNet(nn.Module):\n",
    "    def __init__(self, in_channels=3, out_channels=3, n_residual_blocks=16):\n",
    "        super(GeneratorResNet, self).__init__()\n",
    "\n",
    "        # First layer\n",
    "        self.conv1 = nn.Sequential(nn.Conv2d(in_channels, 64, kernel_size=9, stride=1, padding=4), nn.LeakyReLU())\n",
    "\n",
    "        # Residual blocks\n",
    "        res_blocks = []\n",
    "        for _ in range(n_residual_blocks):\n",
    "            res_blocks.append(ResidualBlock(64))\n",
    "        self.res_blocks = nn.Sequential(*res_blocks)\n",
    "\n",
    "        # Second conv layer post residual blocks\n",
    "        self.conv2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1), nn.BatchNorm2d(64, 0.8))\n",
    "\n",
    "        # Upsampling layers\n",
    "        upsampling = []\n",
    "        for out_features in range(2):\n",
    "            upsampling += [\n",
    "                # nn.Upsample(scale_factor=2),\n",
    "                nn.Conv2d(64, 256, 3, 1, 1),\n",
    "                nn.BatchNorm2d(256),\n",
    "                nn.PixelShuffle(upscale_factor=2),\n",
    "                nn.PReLU(),\n",
    "            ]\n",
    "        self.upsampling = nn.Sequential(*upsampling)\n",
    "\n",
    "        # Final output layer\n",
    "        self.conv3 = nn.Sequential(nn.Conv2d(64, out_channels, kernel_size=9, stride=1, padding=4), nn.Tanh())\n",
    "\n",
    "    def forward(self, x):\n",
    "        out1 = self.conv1(x)\n",
    "        out = self.res_blocks(out1)\n",
    "        out2 = self.conv2(out)\n",
    "        out = torch.add(out1, 0.85*out2)\n",
    "        out = self.upsampling(out)\n",
    "        out = self.conv3(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_shape):\n",
    "        super(Discriminator, self).__init__()\n",
    "\n",
    "        self.input_shape = input_shape\n",
    "        in_channels, in_height, in_width = self.input_shape\n",
    "        patch_h, patch_w = int(in_height / 2 ** 4), int(in_width / 2 ** 4)\n",
    "        self.output_shape = (1, patch_h, patch_w)\n",
    "\n",
    "        def discriminator_block(in_filters, out_filters, first_block=False):\n",
    "            layers = []\n",
    "            layers.append(nn.Conv2d(in_filters, out_filters, kernel_size=3, stride=1, padding=1))\n",
    "            if not first_block:\n",
    "                layers.append(nn.BatchNorm2d(out_filters))\n",
    "            layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
    "            layers.append(nn.Conv2d(out_filters, out_filters, kernel_size=3, stride=2, padding=1))\n",
    "            layers.append(nn.BatchNorm2d(out_filters))\n",
    "            layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
    "            return layers\n",
    "\n",
    "        layers = []\n",
    "        in_filters = in_channels\n",
    "        for i, out_filters in enumerate([64, 128, 256, 512]):\n",
    "            layers.extend(discriminator_block(in_filters, out_filters, first_block=(i == 0)))\n",
    "            in_filters = out_filters\n",
    "\n",
    "        layers.append(nn.Conv2d(out_filters, 1, kernel_size=3, stride=1, padding=1))\n",
    "\n",
    "        self.model = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, img):\n",
    "        return self.model(img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "2**4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64\n",
      "128\n",
      "256\n",
      "512\n"
     ]
    }
   ],
   "source": [
    "for i, out_filters in enumerate([64, 128, 256, 512]):\n",
    "    print(out_filters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train Super Resolution GAN (SRGAN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize generator and discriminator\n",
    "generator = GeneratorResNet()\n",
    "discriminator = Discriminator(input_shape=(channels, *hr_shape))\n",
    "feature_extractor = FeatureExtractor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set feature extractor to inference mode\n",
    "feature_extractor.eval()\n",
    "\n",
    "# Losses\n",
    "criterion_GAN = torch.nn.MSELoss()\n",
    "criterion_GAN1=torch.nn.BCEWithLogitsLoss(reduction='mean')\n",
    "criterion_content = torch.nn.L1Loss()\n",
    "\n",
    "device=torch.device(\"mps\")\n",
    "\n",
    "generator.to(device)\n",
    "discriminator.to(device)\n",
    "feature_extractor.to(device)\n",
    "criterion_GAN.to(device)\n",
    "criterion_content.to(device)\n",
    "\n",
    "# Load pretrained models\n",
    "# if load_pretrained_models:\n",
    "#     generator.load_state_dict(torch.load(\"../input/single-image-super-resolution-gan-srgan-pytorch/saved_models/generator.pth\"))\n",
    "#     discriminator.load_state_dict(torch.load(\"../input/single-image-super-resolution-gan-srgan-pytorch/saved_models/discriminator.pth\"))\n",
    "\n",
    "# Optimizers\n",
    "optimizer_G = torch.optim.Adam(generator.parameters(), lr=lr, betas=(b1, b2))\n",
    "optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(b1, b2))\n",
    "\n",
    "Tensor = torch.Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_gen_losses, train_disc_losses, train_counter = [], [], []\n",
    "test_gen_losses, test_disc_losses = [], []\n",
    "test_counter = [idx * len(train_dataloader.dataset) for idx in range(1, n_epochs + 1)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e896ec2a69df4484ba7e6a18a3a4c8c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training Epoch 0 :   0%|          | 0/13 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b40b11e602d4be29b0b9438445b5bfa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing Epoch 0 :   0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "ename": "TypeError",
     "evalue": "forward() missing 1 required positional argument: 'x'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[18], line 112\u001b[0m\n\u001b[1;32m    110\u001b[0m \u001b[38;5;66;03m# Save model checkpoints\u001b[39;00m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m np\u001b[38;5;241m.\u001b[39margmin(test_gen_losses) \u001b[38;5;241m==\u001b[39m \u001b[38;5;28mlen\u001b[39m(test_gen_losses) \u001b[38;5;241m-\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m--> 112\u001b[0m     torch\u001b[38;5;241m.\u001b[39msave(\u001b[43mgenerator\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msaved_models/generator.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    113\u001b[0m     torch\u001b[38;5;241m.\u001b[39msave(discriminator(), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msaved_models/discriminator.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/torch-gpu/lib/python3.9/site-packages/torch/nn/modules/module.py:1532\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1530\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1531\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1532\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/torch-gpu/lib/python3.9/site-packages/torch/nn/modules/module.py:1541\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1536\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1537\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1538\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1539\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1540\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1541\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1543\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1544\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "\u001b[0;31mTypeError\u001b[0m: forward() missing 1 required positional argument: 'x'"
     ]
    }
   ],
   "source": [
    "    for epoch in range(5):\n",
    "\n",
    "        ### Training\n",
    "        gen_loss, disc_loss = 0, 0\n",
    "        tqdm_bar = tqdm(train_dataloader, desc=f'Training Epoch {epoch} ', total=int(len(train_dataloader)))\n",
    "        for batch_idx, imgs in enumerate(tqdm_bar):\n",
    "            generator.train();\n",
    "            discriminator.train()\n",
    "            # Configure model input\n",
    "            imgs_lr = Variable(imgs[\"lr\"].type(Tensor))\n",
    "            imgs_hr = Variable(imgs[\"hr\"].type(Tensor))\n",
    "            imgs_lr=imgs_lr.to(device)\n",
    "            imgs_hr=imgs_hr.to(device)\n",
    "            # Adversarial ground truths\n",
    "            valid = Variable(Tensor(np.ones((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "            fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "\n",
    "            ### Train Generator\n",
    "            optimizer_G.zero_grad()\n",
    "            # Generate a high resolution image from low resolution input\n",
    "            gen_hr = generator(imgs_lr)\n",
    "            gen_hr=gen_hr.to(device)\n",
    "            # Adversarial loss\n",
    "            valid=valid.to(device)\n",
    "            fake=fake.to(device)\n",
    "            loss_GAN = criterion_GAN(discriminator(gen_hr), valid)\n",
    "            # Content loss\n",
    "            gen_features = feature_extractor(gen_hr)\n",
    "            real_features = feature_extractor(imgs_hr)\n",
    "            loss_content = criterion_content(gen_features, real_features.detach())\n",
    "            # Total loss\n",
    "            loss_G = loss_content + 1e-3 * loss_GAN\n",
    "            loss_G=loss_G.to(device)\n",
    "            loss_G.backward()\n",
    "            optimizer_G.step()\n",
    "\n",
    "            ### Train Discriminator\n",
    "            optimizer_D.zero_grad()\n",
    "            # Loss of real and fake images\n",
    "            loss_real = criterion_GAN(discriminator(imgs_hr), valid)\n",
    "            loss_fake = criterion_GAN(discriminator(gen_hr.detach()), fake)\n",
    "            # Total loss\n",
    "            loss_D = (loss_real + loss_fake) / 2\n",
    "            loss_D.backward()\n",
    "            optimizer_D.step()\n",
    "\n",
    "            gen_loss += loss_G.item()\n",
    "            train_gen_losses.append(loss_G.item())\n",
    "            disc_loss += loss_D.item()\n",
    "            train_disc_losses.append(loss_D.item())\n",
    "            train_counter.append(batch_idx * batch_size + imgs_lr.size(0) + epoch * len(train_dataloader.dataset))\n",
    "            tqdm_bar.set_postfix(gen_loss=gen_loss / (batch_idx + 1), disc_loss=disc_loss / (batch_idx + 1))\n",
    "        # Testing\n",
    "        gen_loss, disc_loss = 0, 0\n",
    "        tqdm_bar = tqdm(test_dataloader, desc=f'Testing Epoch {epoch} ', total=int(len(test_dataloader)))\n",
    "        for batch_idx, imgs in enumerate(tqdm_bar):\n",
    "            generator.eval();\n",
    "            discriminator.eval()\n",
    "            # Configure model input\n",
    "            imgs_lr = Variable(imgs[\"lr\"].type(Tensor))\n",
    "            imgs_hr = Variable(imgs[\"hr\"].type(Tensor))\n",
    "            # Adversarial ground truths\n",
    "            imgs_lr = imgs_lr.to(device)\n",
    "            imgs_hr = imgs_hr.to(device)\n",
    "\n",
    "            valid = Variable(Tensor(np.ones((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "            fake = Variable(Tensor(np.zeros((imgs_lr.size(0), *discriminator.output_shape))), requires_grad=False)\n",
    "            valid=valid.to(device)\n",
    "            fake=fake.to(device)\n",
    "            ### Eval Generator\n",
    "            # Generate a high resolution image from low resolution input\n",
    "            gen_hr = generator(imgs_lr)\n",
    "            # Adversarial loss\n",
    "            gen_hr=gen_hr.to(device)\n",
    "            loss_GAN = criterion_GAN(discriminator(gen_hr), valid)\n",
    "            # Content loss\n",
    "            loss_GAN=loss_GAN.to(device)\n",
    "            gen_features = feature_extractor(gen_hr)\n",
    "            real_features = feature_extractor(imgs_hr)\n",
    "            gen_features=gen_features.to(device)\n",
    "            real_features=real_features.to(device)\n",
    "            loss_content = criterion_content(gen_features, real_features.detach())\n",
    "            # Total loss\n",
    "\n",
    "            loss_G = loss_content + 1e-3 * loss_GAN\n",
    "\n",
    "            ### Eval Discriminator\n",
    "            # Loss of real and fake images\n",
    "            loss_real = criterion_GAN(discriminator(imgs_hr), valid)\n",
    "            loss_fake = criterion_GAN(discriminator(gen_hr.detach()), fake)\n",
    "            # Total loss\n",
    "            loss_D = (loss_real + loss_fake) / 2\n",
    "\n",
    "            gen_loss += loss_G.item()\n",
    "            disc_loss += loss_D.item()\n",
    "            tqdm_bar.set_postfix(gen_loss=gen_loss / (batch_idx + 1), disc_loss=disc_loss / (batch_idx + 1))\n",
    "\n",
    "            # Save image grid with upsampled inputs and SRGAN outputs\n",
    "            if random.uniform(0, 1) < 0.1:\n",
    "                imgs_lr = nn.functional.interpolate(imgs_lr, scale_factor=4)\n",
    "                imgs_hr = make_grid(imgs_hr, nrow=1, normalize=True)\n",
    "                gen_hr = make_grid(gen_hr, nrow=1, normalize=True)\n",
    "                imgs_lr = make_grid(imgs_lr, nrow=1, normalize=True)\n",
    "                img_grid = torch.cat((imgs_hr, imgs_lr, gen_hr), -1)\n",
    "                save_image(img_grid, f\"images/{batch_idx}.png\", normalize=False)\n",
    "\n",
    "        test_gen_losses.append(gen_loss / len(test_dataloader))\n",
    "        test_disc_losses.append(disc_loss / len(test_dataloader))\n",
    "\n",
    "        # Save model checkpoints\n",
    "        if np.argmin(test_gen_losses) == len(test_gen_losses) - 1:\n",
    "            torch.save(generator, \"saved_models/generator.pth\")\n",
    "            torch.save(discriminator, \"saved_models/discriminator.pth\")\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "fig.add_trace(go.Scatter(x=train_counter, y=train_gen_losses, mode='lines', name='Train Generator Loss'))\n",
    "fig.add_trace(go.Scatter(x=test_counter, y=test_gen_losses, marker_symbol='star-diamond', \n",
    "                         marker_color='orange', marker_line_width=1, marker_size=9, mode='markers', name='Test Generator Loss'))\n",
    "fig.update_layout(\n",
    "    width=1000,\n",
    "    height=500,\n",
    "    title=\"Train vs. Test Generator Loss\",\n",
    "    xaxis_title=\"Number of training examples seen\",\n",
    "    yaxis_title=\"Adversarial + Content Loss\"),\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "fig.add_trace(go.Scatter(x=train_counter, y=train_disc_losses, mode='lines', name='Train Discriminator Loss'))\n",
    "fig.add_trace(go.Scatter(x=test_counter, y=test_disc_losses, marker_symbol='star-diamond', \n",
    "                         marker_color='orange', marker_line_width=1, marker_size=9, mode='markers', name='Test Discriminator Loss'))\n",
    "fig.update_layout(\n",
    "    width=1000,\n",
    "    height=500,\n",
    "    title=\"Train vs. Test Discriminator Loss\",\n",
    "    xaxis_title=\"Number of training examples seen\",\n",
    "    yaxis_title=\"Adversarial Loss\"),\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"test\", exist_ok=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_model=torch.load(\"saved_model/generator.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img=Image.open(\"/kaggle/input/celeba-dataset/img_align_celeba/img_align_celeba/000001.jpg\")\n",
    "trans=transforms.Compose([transforms.Resize((64,64)),transforms.ToTensor()])\n",
    "img=trans(img)\n",
    "img=torch.reshape(img,(1,3,64,64))\n",
    "img=img.cuda()\n",
    "gen_hr = test_model(img)\n",
    "save_image(gen_hr, \"test/test.png\", normalize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img=Image.open(\"/kaggle/input/celeba-dataset/img_align_celeba/img_align_celeba/000001.jpg\")\n",
    "trans=transforms.Compose([transforms.Resize((64,64)),transforms.ToTensor()])\n",
    "img=trans(img)\n",
    "img=torch.reshape(img,(1,3,64,64))\n",
    "img=img.cuda()\n",
    "gen_hr = generator(img)\n",
    "\n",
    "save_image(gen_hr, \"test/test.png\", normalize=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img=Image.open(\"/kaggle/working/test/test.png\")\n",
    "trans=transforms.Compose([transforms.Resize((178,218)),transforms.ToTensor()])\n",
    "img=trans(img)\n",
    "img=torch.reshape(img,(1,3,178,218))\n",
    "img\n",
    "save_image(img, \"test/test1.png\", normalize=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Work in Progress ...\n",
    "### Note: The order of the below saved images are (from left-to-right): Original HR, Bicubic Interpolated & SRGAN Generated"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "datasetId": 29561,
     "sourceId": 37705,
     "sourceType": "datasetVersion"
    },
    {
     "sourceId": 51918360,
     "sourceType": "kernelVersion"
    }
   ],
   "dockerImageVersionId": 30018,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
